import os
import json
import numpy as np
import scipy.io as scio
import hdf5storage as hdf5

def get_info(src_dir):

    names = os.listdir(src_dir)
    names_sel, types = [], []

    for name in names:
        mal_path = src_dir + name + '\\'
        fnames = os.listdir(mal_path)
        for fname in fnames:
            flag = 0
            if '.json' in fname:
                flag = 1
                fjson = mal_path + fname
                # parse json
                try:
                    f = open(fjson, 'r')
                    cjson = json.load(f)
                    mtype = str(cjson['scans']['Avast']['result'])
                except:
                    flag = 0
            if flag:
                types.append(mtype)
                names_sel.append(name)

    f = open('info.csv', 'w')
    for it in range(len(types)):
        tmp = names_sel[it]+','+types[it]+'\n'
        f.write(tmp)
    f.close()

    names_nor = []
    for name in names:
        if 'Normal' in name:
            nor_path = src_dir + name + '\\'
            fnames = os.listdir(nor_path)
            for fname in fnames:
                flag = 0
                if '.txt' in fname:
                    names_nor.append(name)

    for name in names_nor:
        print(name)

def fuse_into_txt(src_dir):

    # maltypes = ['Artemis', 'BackDoor', 'Downloader', 'Drixed', 'Emotet', 'Fareit', 'Generic', 'Injector', 'JS',
    #             'Medfos', 'None', 'Obfuscated', 'Packed', 'PUP', 'PWS', 'Ransom', 'RemAdm', 'Trojan', 'Upatre',
    #             'Netsky', 'Ramnit', 'Sality', 'Sdbot', 'Virut', 'ZeroAccess', 'Normal']
    maltypes = ['Adw', 'Drp', 'Rtk', 'Susp', 'Trj', 'PUP', 'Expl', 'Cryp', 'Wrm', 'Normal']

    # get the infos
    info_path = 'info.csv'
    f = open(info_path, 'r')
    names, labels, nums = [], [], []
    for line in f:
        tmp = line[:-1].split(',')
        names.append(tmp[0].strip())
        labels.append(tmp[1].strip())
        # nums.append(int(tmp[2].strip()))
    f.close()

    for mt in maltypes:
        fw = open(mt + '.txt', 'w')
        num = 0
        for i in range(len(names)):
            name = names[i]
            if (mt in labels[i]):
                mal_dir = src_dir + name + '\\'
                fnames = os.listdir(mal_dir)
                for fname in fnames:
                    if '.txt' in fname:
                        fea_path = mal_dir + fname
                        ff = open(fea_path, 'r')
                        for line in ff:
                            fw.write(line[:-3] + '\n')
                            num += 1
                        ff.close()
        fw.close()
        print(mt, ': ', num)

def convert_txt_into_npz(src_dir):

    maltypes = ['Normal', 'Adw', 'Drp', 'Rtk', 'Susp', 'Trj']
    for im in range(len(maltypes)):
        mt = maltypes[im]
        print('- ', mt)

        txt_path = src_dir + mt + '.txt'
        f = open(txt_path, 'r')

        packet, tls, cert = [], [], []
        # Xs = [np.zeros([1,27]), np.zeros([1,786]), np.zeros([1,42])]
        num = 0

        for line in f:
            tmp = line.split(' ')
            ta = [float(t) for t in tmp]
            # Xs[0] = np.vstack((Xs[0], np.array(ta[0:27])))
            # Xs[1] = np.vstack((Xs[1], np.array(ta[27:813])))
            # Xs[2] = np.vstack((Xs[2], np.array(ta[813:855])))
            packet.append(ta[0:27])
            tls.append(ta[27:813])
            cert.append(ta[813:855])
            num += 1
        f.close()

        packet = np.array(packet)
        tls = np.array(tls)
        cert = np.array(cert)
        gt = np.ones(num, dtype=np.int64) * im

        data_name = 'CTU_Encrypted_Malware_Traffic'
        type = mt

        save_path = src_dir + mt + '.npz'
        np.savez(save_path, packet=packet, tls=tls, cert=cert, gt=gt, data_name=data_name, type=type)

def convert_npz_into_mat(data_dir):

    names = ['Normal', 'Adw', 'Drp', 'Rtk', 'Susp', 'Trj']

    class_meaning = np.array(names, dtype=object)
    X, Y = [[], [], []], []
    for name in names:
        data_path = data_dir + name + '.npz'
        data = np.load(data_path)
        X[0].append(data['packet'])
        X[1].append(data['tls'])
        X[2].append(data['cert'])
        Y.append(data['gt'])
    X[0] = np.vstack(X[0]).T
    X[1] = np.vstack(X[1]).T
    X[2] = np.vstack(X[2]).T
    X = np.array(X, dtype=object).T
    Y = np.hstack(Y).T

    data_name = 'CTU_Encrypted_Malware_Traffic'
    view_meaning = np.array(['packet', 'tls', 'cert'], dtype=object)

    scio.savemat(data_dir + data_name + '.mat', {'X': X, 'Y': Y, 'data_name': data_name, 'class_meaning': class_meaning,
                                                 'view_meaning': view_meaning})
    hdf5.write({'X': X, 'Y': Y, 'data_name': data_name, 'class_meaning': class_meaning, 'view_meaning': view_meaning}, path='./', filename=data_name+'.mat', truncate_invalid_matlab=True)

    return

def gen_single_mal_mat(data_dir):

    names = ['Normal', 'Adw', 'Drp', 'Rtk', 'Susp', 'Trj']

    for i in range(1,len(names)):
        class_meaning = np.array(['Normal', names[i]], dtype=object)
        X, Y = [[], [], []], []
        for name in ['Normal', names[i]]:
            data_path = data_dir + name + '.npz'
            data = np.load(data_path)
            X[0].append(data['packet'])
            X[1].append(data['tls'])
            X[2].append(data['cert'])
            Y.append(data['gt'])
        X[0] = np.vstack(X[0]).T
        X[1] = np.vstack(X[1]).T
        X[2] = np.vstack(X[2]).T
        X = np.array(X, dtype=object).T
        Y = np.hstack(Y).T
        Y[np.argwhere(Y!=0)] = 1

        data_name = 'CTU_Normal_vs_{}'.format(names[i])
        view_meaning = np.array(['packet', 'tls', 'cert'], dtype=object)
        hdf5.write(
            {'X': X, 'Y': Y, 'data_name': data_name, 'class_meaning': class_meaning, 'view_meaning': view_meaning},
            path='./', filename=data_name + '.mat', truncate_invalid_matlab=True)

    class_meaning = np.array(names[:-1], dtype=object)
    X, Y = [[], [], []], []
    for name in names[:-1]:
        data_path = data_dir + name + '.npz'
        data = np.load(data_path)
        X[0].append(data['packet'])
        X[1].append(data['tls'])
        X[2].append(data['cert'])
        Y.append(data['gt'])
    X[0] = np.vstack(X[0]).T
    X[1] = np.vstack(X[1]).T
    X[2] = np.vstack(X[2]).T
    X = np.array(X, dtype=object).T
    Y = np.hstack(Y).T

    data_name = 'CTU_Normal_vs_Adw_Drp_Rtk_Susp'
    view_meaning = np.array(['packet', 'tls', 'cert'], dtype=object)
    hdf5.write(
        {'X': X, 'Y': Y, 'data_name': data_name, 'class_meaning': class_meaning, 'view_meaning': view_meaning},
        path='./', filename=data_name + '.mat', truncate_invalid_matlab=True)

    return

if __name__ == '__main__':

    # src_dir = 'D:\\Work\\datasets\\Dataset\\'

    # get_info(src_dir)

    # fuse_into_txt(src_dir)

    # convert_txt_into_npz('./')
    # convert_npz_into_mat('./')

    gen_single_mal_mat('./')


